import os
import numpy as np
import torch
import shutil
import torchvision.transforms as transforms
from torch.autograd import Variable



class AvgrageMeter(object):

  def __init__(self):
    self.reset()

  def reset(self):
    self.avg = 0
    self.sum = 0
    self.cnt = 0

  def update(self, val, n=1):
    self.sum += val * n
    self.cnt += n
    self.avg = self.sum / self.cnt


def accuracy(output, target, topk=(1,)):
  maxk = max(topk)
  batch_size = target.size(0)

  _, pred = output.topk(maxk, 1, True, True)
  pred = pred.t()
  correct = pred.eq(target.view(1, -1).expand_as(pred))

  res = []
  for k in topk:
    correct_k = correct[:k].contiguous().view(-1).float().sum(0)
    res.append(correct_k.mul_(100.0/batch_size))
  return res


class Cutout(object):
    def __init__(self, length):
        self.length = length

    def __call__(self, img):
        h, w = img.size(1), img.size(2)
        mask = np.ones((h, w), np.float32)
        y = np.random.randint(h)
        x = np.random.randint(w)

        y1 = np.clip(y - self.length // 2, 0, h)
        y2 = np.clip(y + self.length // 2, 0, h)
        x1 = np.clip(x - self.length // 2, 0, w)
        x2 = np.clip(x + self.length // 2, 0, w)

        mask[y1: y2, x1: x2] = 0.
        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img *= mask
        return img

# def _data_transforms_ham10000(args):
#     """
#     Generate transformations for the HAM10000 dataset based on provided args.
#
#     Args:
#     - args: Configuration object with attributes such as resize, rotation_degree, mean, std, use_cutout, cutout_length.
#
#     Returns:
#     - train_transform: The transformation for training data.
#     - valid_transform: The transformation for validation data.
#     """
#     # Default values are set if not specified in args
#     resize = getattr(args, 'resize', (224, 224))
#     rotation_degree = getattr(args, 'rotation_degree', 20)
#     mean = getattr(args, 'mean', [0.763038, 0.545646, 0.570039])
#     std = getattr(args, 'std', [0.140928, 0.152613, 0.169973])
#     use_cutout = getattr(args, 'use_cutout', False)
#     cutout_length = getattr(args, 'cutout_length', 16)
#
#     train_transform_list = [
#         transforms.Resize(resize),
#         transforms.RandomHorizontalFlip(),
#         transforms.RandomRotation(rotation_degree),
#         transforms.ToTensor(),
#         transforms.Normalize(mean, std)
#     ]
#
#     if use_cutout:
#         # Assuming Cutout is a custom transformation class you have defined or imported
#         train_transform_list.append(Cutout(cutout_length))
#
#     train_transform = transforms.Compose(train_transform_list)
#
#     valid_transform = transforms.Compose([
#         transforms.Resize(resize),
#         transforms.ToTensor(),
#         transforms.Normalize(mean, std)
#     ])
#
#     return train_transform, valid_transform
from imgaug import augmenters as iaa
from PIL import Image
class ImgAugTransform:
    def __init__(self):
        self.aug = iaa.Sequential([
            iaa.ElasticTransformation(alpha=50, sigma=5)  # 示例，可以根据需求更改
        ])

    def __call__(self, img):
        img = np.array(img)  # 将 PIL Image 转换为 numpy array
        img = self.aug(image=img)  # 应用 imgaug 处理
        img = Image.fromarray(img.astype('uint8'), 'RGB')  # 将 numpy array 转回 PIL Image
        return img
def _data_transforms_ham10000(args):
    # 设置参数
    resize = getattr(args, 'resize', (192, 192))
    rotation_degree = getattr(args, 'rotation_degree', 20)
    mean = getattr(args, 'mean', [0.763038, 0.545646, 0.570039])
    std = getattr(args, 'std', [0.140928, 0.152613, 0.169973])
    use_cutout = getattr(args, 'use_cutout', False)
    cutout_length = getattr(args, 'cutout_length', 16)
    train_transform_list = [
        ImgAugTransform(),  # 弹性形变
        transforms.Resize(resize),
        transforms.RandomCrop(64),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(rotation_degree),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ]

    if use_cutout:
        # 确保 Cutout 是一个正确定义的变换，可能需要自定义实现
        train_transform_list.append(Cutout(cutout_length))

    train_transform = transforms.Compose(train_transform_list)

    valid_transform = transforms.Compose([
        transforms.Resize(resize),
        transforms.CenterCrop(64),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    return train_transform, valid_transform
def _data_transforms_cifar10(args):
  CIFAR_MEAN = [0.5374551, 0.55888027, 0.5680246]
  CIFAR_STD = [0.24409676, 0.23635243, 0.2516532]
  Mean=[0.4797, 0.5108, 0.5341]
  Std=[0.1859, 0.1844, 0.2074]

  train_transform = transforms.Compose([
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(Mean, Std),
  ])
  if args.cutout:
    train_transform.transforms.append(Cutout(args.cutout_length))

  valid_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(Mean, Std),
    ])
  return train_transform, valid_transform


def count_parameters_in_MB(model):
  return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name)/1e6


def save_checkpoint(state, is_best, save):
  filename = os.path.join(save, 'checkpoint.pth.tar')
  torch.save(state, filename)
  if is_best:
    best_filename = os.path.join(save, 'model_best.pth.tar')
    shutil.copyfile(filename, best_filename)


def save(model, model_path):
  torch.save(model.state_dict(), model_path)


def load(model, model_path):
  model.load_state_dict(torch.load(model_path))


def drop_path(x, drop_prob):
  if drop_prob > 0.:
    keep_prob = 1.-drop_prob
    mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
    x.div_(keep_prob)
    x.mul_(mask)
  return x


def create_exp_dir(path, scripts_to_save=None):
  if not os.path.exists(path):
    os.mkdir(path)
  print('Experiment dir : {}'.format(path))

  if scripts_to_save is not None:
    os.mkdir(os.path.join(path, 'scripts'))
    for script in scripts_to_save:
      dst_file = os.path.join(path, 'scripts', os.path.basename(script))
      shutil.copyfile(script, dst_file)

